# 人工智能NLP-Agent数字人项目-04-基金数据问答任务工单V1.1-2025.2.14
import sqlite3
import pandas as pd
import utils.configFinRAG as configFinRAG
from FinSQL_01_generate import generate_sql
from FinSQL_02_query import query_db
from FinSQL_03_answer_from_SQL import generate_answer
from utils.instances import TOKENIZER, LLM


# 全局变量初始化
g_example_question_list = []
g_example_sql_list = []
g_example_info_list = []
g_example_fa_list = []
g_example_token_list = []


def load_example_data():
    """加载示例数据到全局变量"""
    if not g_example_question_list:
        sql_examples_file = pd.read_csv(configFinRAG.sql_examples_path, delimiter=",", header=0)
        for _, row in sql_examples_file.iterrows():
            g_example_question_list.append(row['问题'])
            g_example_sql_list.append(row['SQL'])
            g_example_info_list.append(row['资料'])
            g_example_fa_list.append(row['FA'])
            tokens = TOKENIZER(row['问题'])['input_ids']
            g_example_token_list.append(tokens)


def sql_retrieve_chain(query):
    """SQL检索链路的主函数"""
    try:
        # 确保示例数据已加载
        load_example_data()

        # 生成SQL语句
        result_prompt, sql = generate_sql(query, LLM, g_example_question_list, g_example_sql_list, g_example_token_list)

        # 查询数据库
        with sqlite3.connect(configFinRAG.database_path) as conn:
            cursor = conn.cursor()
            success_flag, exc_result = query_db(sql, cursor)
            if not success_flag:
                raise Exception(f"Database query failed: {exc_result}")

        # 生成答案
        answer = generate_answer(query, exc_result, LLM, g_example_question_list, g_example_info_list, g_example_fa_list, g_example_token_list)
        return answer

    except Exception as e:
        print(f"Error in sql_retrieve_chain: {e}")
        return None


if __name__ == "__main__":
    # 示例调用
    query = "示例问题"
    answer = sql_retrieve_chain(query)
    print("Answer:", answer)